import numpy as np
import pandas as pd

# 假设已经读取了处理后的数据列（根据你数据的列名进行处理）
columns = [
    "duration", "*protocol_type_icmp", "*protocol_type_tcp", "*protocol_type_udp",
    "*service_IRC", "*service_X11", "*service_Z39_50", "*service_auth", "*service_bgp",
    "*service_courier", "*service_csnet_ns", "*service_ctf", "*service_daytime", "*service_discard",
    "*service_domain", "*service_domain_u", "*service_echo", "*service_eco_i", "*service_ecr_i",
    "*service_efs", "*service_exec", "*service_finger", "*service_ftp", "*service_ftp_data",
    "*service_gopher", "*service_hostnames", "*service_http", "*service_http_443", "*service_imap4",
    "*service_iso_tsap", "*service_klogin", "*service_kshell", "*service_ldap", "*service_link",
    "*service_login", "*service_mtp", "*service_name", "*service_netbios_dgm", "*service_netbios_ns",
    "*service_netbios_ssn", "*service_netstat", "*service_nnsp", "*service_nntp", "*service_ntp_u",
    "*service_other", "*service_pm_dump", "*service_pop_2", "*service_pop_3", "*service_printer",
    "*service_private", "*service_red_i", "*service_remote_job", "*service_rje", "*service_shell",
    "*service_smtp", "*service_sql_net", "*service_ssh", "*service_sunrpc", "*service_supdup",
    "*service_systat", "*service_telnet", "*service_tftp_u", "*service_tim_i", "*service_time",
    "*service_urh_i", "*service_urp_i", "*service_uucp", "*service_uucp_path", "*service_vmnet",
    "*service_whois", "*flag_OTH", "*flag_REJ", "*flag_RSTO", "*flag_RSTOS0", "*flag_RSTR",
    "*flag_S0", "*flag_S1", "*flag_S2", "*flag_S3", "*flag_SF", "*flag_SH",
    "src_bytes", "dst_bytes", "land", "wrong_fragment", "urgent", "hot", "num_failed_logins",
    "logged_in", "num_compromised", "root_shell", "su_attempted", "num_root",
    "num_file_creations", "num_shells", "num_access_files", "num_outbound_cmds",
    "is_host_login", "is_guest_login",
    "count", "srv_count", "serror_rate", "srv_serror_rate", "rerror_rate",
    "srv_rerror_rate", "same_srv_rate", "diff_srv_rate", "srv_diff_host_rate",
    "dst_host_count", "dst_host_srv_count", "dst_host_same_srv_rate", "dst_host_diff_srv_rate",
    "dst_host_same_src_port_rate", "dst_host_srv_diff_host_rate", "dst_host_serror_rate",
    "dst_host_srv_serror_rate", "dst_host_rerror_rate", "dst_host_srv_rerror_rate",
    "*label_Dos", "*label_Probe", "*label_R2L", "*label_U2R", "*label_normal"
]

# 获取数据列的总数
total_columns = len(columns)

# **特征分类**
# 假设我们已经将每个特征按类别分组
# TCP 连接的基本特征: 包含协议类型、服务类型、标志等
basic_features = [
    "duration", "*protocol_type_icmp", "*protocol_type_tcp", "*protocol_type_udp",
    "*service_IRC", "*service_X11", "*service_Z39_50", "*service_auth", "*service_bgp",
    "*service_courier", "*service_csnet_ns", "*service_ctf", "*service_daytime", "*service_discard",
    "*service_domain", "*service_domain_u", "*service_echo", "*service_eco_i", "*service_ecr_i",
    "*service_efs", "*service_exec", "*service_finger", "*service_ftp", "*service_ftp_data",
    "*service_gopher", "*service_hostnames", "*service_http", "*service_http_443", "*service_imap4",
    "*service_iso_tsap", "*service_klogin", "*service_kshell", "*service_ldap", "*service_link",
    "*service_login", "*service_mtp", "*service_name", "*service_netbios_dgm", "*service_netbios_ns",
    "*service_netbios_ssn", "*service_netstat", "*service_nnsp", "*service_nntp", "*service_ntp_u",
    "*service_other", "*service_pm_dump", "*service_pop_2", "*service_pop_3", "*service_printer",
    "*service_private", "*service_red_i", "*service_remote_job", "*service_rje", "*service_shell",
    "*service_smtp", "*service_sql_net", "*service_ssh", "*service_sunrpc", "*service_supdup",
    "*service_systat", "*service_telnet", "*service_tftp_u", "*service_tim_i", "*service_time",
    "*service_urh_i", "*service_urp_i", "*service_uucp", "*service_uucp_path", "*service_vmnet",
    "*service_whois", "*flag_OTH", "*flag_REJ", "*flag_RSTO", "*flag_RSTOS0", "*flag_RSTR",
    "*flag_S0", "*flag_S1", "*flag_S2", "*flag_S3", "*flag_SF", "*flag_SH"
]

# TCP 连接的内容特征
content_features = [
    "hot", "num_failed_logins", "logged_in", "num_compromised", "root_shell", "su_attempted", 
    "num_root", "num_file_creations", "num_shells", "num_access_files", "num_outbound_cmds", 
    "is_host_login", "is_guest_login"
]

# 基于时间的特征
time_based_features = [
    "count", "srv_count", "serror_rate", "srv_serror_rate", "rerror_rate", "srv_rerror_rate", 
    "same_srv_rate", "diff_srv_rate", "srv_diff_host_rate"
]

# 基于主机的特征
host_based_features = [
    "dst_host_count", "dst_host_srv_count", "dst_host_same_srv_rate", "dst_host_diff_srv_rate", 
    "dst_host_same_src_port_rate", "dst_host_srv_diff_host_rate", "dst_host_serror_rate", 
    "dst_host_srv_serror_rate", "dst_host_rerror_rate", "dst_host_srv_rerror_rate"
]

# 网络攻击类别标签
labels = [
    "*label_Dos", "*label_Probe", "*label_R2L", "*label_U2R", "*label_normal"
]

# **为每个类别生成掩模**
mask_DoS = np.zeros(total_columns, dtype=int)
mask_Probe = np.zeros(total_columns, dtype=int)
mask_U2R_R2L = np.zeros(total_columns, dtype=int)

# **DoS 掩模**: 修改内容特征、基于主机的特征
for feature in content_features + host_based_features:
    mask_DoS[columns.index(feature)] = 1

# **Probe 掩模**: 仅修改内容特征
for feature in content_features:
    mask_Probe[columns.index(feature)] = 1

# **U2R/R2L 掩模**: 修改基于时间的特征、基于主机的特征
for feature in time_based_features + host_based_features:
    mask_U2R_R2L[columns.index(feature)] = 1
mask_DoS[columns.index("*label_Dos")] = 1    
mask_Probe[columns.index("*label_Probe")] = 1
mask_U2R_R2L[columns.index( "*label_R2L")] = 1
mask_U2R_R2L[columns.index( "*label_U2R")] = 1
# 打印掩模
print("DoS 掩模:", mask_DoS)
print("Probe 掩模:", mask_Probe)
print("U2R/R2L 掩模:", mask_U2R_R2L)

# 保存掩模为 CSV 文件
mask_df = pd.DataFrame([mask_DoS, mask_Probe, mask_U2R_R2L], columns=columns)
# mask_df.index = ["DoS", "Probe", "U2R/R2L"]
mask_df.to_csv("adjusted_feature_masks.csv",index=False)

# 打印最终的掩模表格
print(mask_df)
